
import uuid

from threading import RLock
from langchain_huggingface import HuggingFaceEmbeddings
from langgraph.store.memory import InMemoryStore
from typing import Literal, Dict, List
from psycopg_pool import ConnectionPool
from langgraph.store.postgres import PostgresStore

from config import Config
from memory.memory_type import MemoryType


class MemoryBase:

    _lock = RLock()
    _instance = None

    def __new__(cls, *args, **kwargs):
        with cls._lock:
            if cls._instance is None:
                embedding_model = HuggingFaceEmbeddings(
                    model_name=Config.EMBEDDING_MODEL,
                    model_kwargs={'device': Config.DEVICE},
                    encode_kwargs={'normalize_embeddings': False}
                )
                cls.index = {"embed": embedding_model, "dims": Config.EMBEDDING_DIM}
                cls.memory_store = None
                cls._instance = super(MemoryBase, cls).__new__(cls)
        return cls._instance

    def __init__(self, memory_store_type: Literal["inmemory", "pgmemory"]):

        if self.memory_store:
            return

        if memory_store_type == "inmemory":
            # 内存记忆
            self.memory_store = InMemoryStore(index=self.index)
        else:
            # postgres数据库记忆
            pool = ConnectionPool(
                conninfo=f"postgresql://{Config.USER_NAME}:{Config.PASSWORD}@{Config.HOST}:{Config.PORT}/{Config.DBNAME}?sslmode=disable",
                max_size=20,
                kwargs={"autocommit": True, "prepare_threshold": 0},
            )
            self.memory_store = PostgresStore(conn=pool, index=self.index)
            self.memory_store.setup()

    def put(self, memory_type: MemoryType, session_id: str=None, key: str=None, value: Dict=None, index: List | None = None) -> None:
        self.memory_store.put(
            namespace = (memory_type.value, session_id) if session_id else (memory_type.value,),
            key=key if key else str(uuid.uuid4()),
            value=value,
            index=index,
        )

    async def aput(self, memory_type: MemoryType, session_id: str=None, key: str=None, value: Dict=None, index: List | None = None) -> None:
        await self.memory_store.aput(
            namespace = (memory_type.value, session_id) if session_id else (memory_type.value,),
            key=key if key else str(uuid.uuid4()),
            value=value,
            index=index,
        )

    def get(self, memory_type: MemoryType, session_id: str = None, key: str = None) -> Dict:
        memory = self.memory_store.get(
            namespace = (memory_type.value, session_id) if session_id else (memory_type.value,),
            key=key
        )
        if memory is None:
            return {}
        return memory.value

    async def aget(self, memory_type: MemoryType, session_id: str, key: str) -> Dict:
        memory = await self.memory_store.aget(
            namespace=(memory_type.value, session_id) if session_id else (memory_type.value,),
            key=key
        )
        if not memory:
            return {}
        return memory.value

    def search(self, memory_type: MemoryType, session_id: str=None, query: str=None, limit: int = 10) -> List[Dict]:
        memories = self.memory_store.search(
            (memory_type.value, session_id, ) if session_id else (memory_type.value,),
            query=query,
            limit=10,
        )
        if not memories:
            return []

        return [memory.value for memory in memories if memory.value][:limit]

    async def asearch(self, memory_type: MemoryType, session_id: str, query: str, limit: int = 10) -> List[Dict]:
        memories = await self.memory_store.asearch(
            (memory_type.value, session_id,) if session_id else (memory_type.value,),
            query=query,
            limit=10,
        )
        if not memories:
            return []

        return [memory.value for memory in memories if memory.value][:limit]
